package utils.fishersExact;

import jsc.contingencytables.ContingencyTable2x2;
import jsc.tests.H1;
import jsc.distributions.Hypergeometric;

public class FishersExactTest {

    private int populationSize;
    private int totalSuccesses;
    private int sampleSize;
    private int sampleSuccesses;
    private ContingencyTable2x2 ct;
    private H1 alternative = H1.NOT_EQUAL;
    public static final int OVERREP = 1;
    public static final int UNDERREP = -1;
    public static final int NOREP = 0;
    private int representation;
    private int z;
    private double foldChange;

    /**
     * A two-tailed Fisher's exact test. Use the getResult() method to get the
     * probability. You will need to have the cern.jet.random package (colt.jar)
     * on your classpath to use this.
     *
     * @param populationSize
     *            The size of the population.
     * @param totalSuccesses
     *            The total number of "successes" that meet the given criteria.
     * @param sampleSize
     *            The size of the sample.
     * @param sampleSuccesses
     *            The total number of "successes" within the sample that meet
     *            the given criteria.
     */
    public FishersExactTest(int populationSize, int totalSuccesses, int sampleSize, int sampleSuccesses) {
        try {
            setVariables(populationSize, totalSuccesses, sampleSize, sampleSuccesses);
        } catch (SillyInputException e) {
            System.out.println(e.toString());
            e.printStackTrace();
        }
    }

    public FishersExactTest(int populationSize, int totalSuccesses, int sampleSize, int sampleSuccesses, H1 alternative) {
        try {
            this.alternative = alternative;
            setVariables(populationSize, totalSuccesses, sampleSize, sampleSuccesses);
        } catch (SillyInputException e) {
            System.out.println(e.toString());
            e.printStackTrace();
        }
    }

    /**
     * Reset the variables for the two-tailed Fisher's exact test. Use the
     * getResult() method to get the probability.
     *
     * @param populationSize
     *            The size of the population
     * @param totalSuccesses
     *            The total number of "successes" that meet the given criteria
     * @param sampleSize
     *            The size of the sample
     * @param sampleSuccesses
     *            The total number of "successes" within the sample that meet
     *            the given criteria
     * @throws SillyInputException
     *             An exception thrown on the input of impossible counts
     */
    public void setVariables(int populationSize, int totalSuccesses, int sampleSize, int sampleSuccesses) throws SillyInputException {
        this.populationSize = populationSize;
        this.totalSuccesses = totalSuccesses;
        this.sampleSize = sampleSize;
        this.sampleSuccesses = sampleSuccesses;
        testInput();
        double expectedProportion = (totalSuccesses + 0.0) / populationSize;
        double expectedSamSuc = sampleSize * expectedProportion;
        foldChange = sampleSuccesses / expectedSamSuc;
        if (expectedSamSuc < sampleSuccesses) {
            representation = OVERREP;
        } else if (expectedSamSuc > sampleSuccesses) {
            representation = UNDERREP;
        } else {
            representation = NOREP;
        }

        int[][] data = new int[2][2];
        data[0][0] = sampleSuccesses;
        data[0][1] = sampleSize - sampleSuccesses;
        data[1][0] = totalSuccesses - sampleSuccesses;
        data[1][1] = populationSize - data[0][0] - data[0][1] - data[1][0];
        try{
            ct = new ContingencyTable2x2(data);
        } catch (IllegalArgumentException ex){
            System.out.println(populationSize+", "+totalSuccesses+", "+sampleSize+", "+sampleSuccesses);
            ex.printStackTrace();
        }

    }

    public double getFoldChange() {
        return foldChange;
    }

    public int getRepresentation() {
        return representation;
    }

    /**
     * Get the results of the test
     *
     * @return The results of the two-tailed Fishers exact test.
     */
    public double getResult() {
        if (sampleSuccesses != 0) {
            if (populationSize - (totalSuccesses - sampleSuccesses) - (sampleSize - sampleSuccesses) - sampleSuccesses == 0) {
                if (alternative == H1.NOT_EQUAL) {
                    try{
                    Hypergeometric hg = new Hypergeometric(sampleSize, populationSize, totalSuccesses);
                    double p0 = hg.pdf(sampleSuccesses);

                    double sum = p0;
                    int max = sampleSize;
                    if (sampleSize > totalSuccesses) {
                        max = totalSuccesses;
                    }

                    if (max >= 1) {
                        for (int i = sampleSuccesses + 1; i <= max; i++) {
                            double pI = hg.pdf(i);
                            if (pI <= p0) {
                                sum += pI;
                            }
                        }
                    }
                    if (alternative == H1.NOT_EQUAL) {
                        return sum;
                    } else if (representation == OVERREP && alternative == H1.GREATER_THAN) {
                        return sum / 2;
                    } else if (representation == UNDERREP && alternative == H1.GREATER_THAN) {
                        return 1 - (sum / 2);
                    } else if (representation == UNDERREP && alternative == H1.LESS_THAN) {
                        return sum / 2;
                    } else if (representation == UNDERREP && alternative == H1.GREATER_THAN) {
                        return 1 - (sum / 2);
                    }
                    return sum;
                    } catch (IllegalArgumentException ex){
                        ex.printStackTrace();
                        return 1.0;
                    }

                }
            } else if (sampleSize != sampleSuccesses && sampleSuccesses != totalSuccesses) {

                jsc.contingencytables.FishersExactTest fet = new jsc.contingencytables.FishersExactTest(ct, alternative);
                return fet.getSP();

            } else {
                if (alternative == H1.GREATER_THAN) {
                    Hypergeometric hg = new Hypergeometric(sampleSize, populationSize, totalSuccesses);
                    return hg.pdf(sampleSuccesses);
                } else if (alternative == H1.LESS_THAN) {
                    return 1.0;
                } else {
                    Hypergeometric hg = new Hypergeometric(sampleSize, populationSize, totalSuccesses);
                    double p0 = hg.pdf(sampleSuccesses);
                    double sum = p0;
                    int max = sampleSize;
                    if (sampleSize > totalSuccesses) {
                        max = totalSuccesses;
                    }
                    int minSuc = 0;
                    if (sampleSize - (populationSize - totalSuccesses) > 0) {
                        minSuc = sampleSize - (populationSize - totalSuccesses);
                    }
                    for (int i = minSuc; i <= max; i++) {
                        double pI = hg.pdf(i);
                        if (pI <= p0) {
                            sum += pI;
                        }
                    }
                    return sum;
                }

            }
        } else if (totalSuccesses == 0) {
            return 1.0;
        } else if (sampleSize == 0) {
            return 1.0;
        } else {
            if (alternative == H1.LESS_THAN) {

                Hypergeometric hg = new Hypergeometric(sampleSize, populationSize, totalSuccesses);
                return hg.pdf(sampleSuccesses);

            } else if (alternative == H1.GREATER_THAN) {
                return 1.0;
            } else {

                Hypergeometric hg = new Hypergeometric(sampleSize, populationSize, totalSuccesses);
                double p0 = hg.pdf(sampleSuccesses);
                double sum = 0;
                int max = sampleSize;
                if (sampleSize > totalSuccesses) {
                    max = totalSuccesses;
                }
                int minSuc = 0;
                if (sampleSize - (populationSize - totalSuccesses) > 0) {
                    minSuc = sampleSize - (populationSize - totalSuccesses);
                }
                for (int i = minSuc; i <= max; i++) {

                    double pI = hg.pdf(i);
                    if (pI <= p0) {
                        sum += pI;
                    }
                }
                return sum;

            }

        }
        return 1.0;

    }

    private void testInput() throws SillyInputException {
        if (populationSize < totalSuccesses) {
            throwException();
        } else if (populationSize < sampleSize) {
            throwException();
        } else if (sampleSize < sampleSuccesses) {
            throwException();
        } else if (totalSuccesses < sampleSuccesses) {
            throwException();
        }
    }

    private void throwException() throws SillyInputException {
        throw new SillyInputException(populationSize, totalSuccesses, sampleSize, sampleSuccesses);
    }
}
